{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%%import tensorflow as tf\n",
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "data": {
      "text/plain": "array([[ 711,  632,   71,    0,    0,    0],\n       [  73,    8, 3215,   55,  927,    0],\n       [  83,   91,    1,  645, 1253,  927]])"
     },
     "metadata": {},
     "output_type": "execute_result",
     "execution_count": 4
    }
   ],
   "source": [
    "\n",
    "raw_inputs = [\n",
    "    [711, 632, 71],\n",
    "    [73, 8, 3215, 55, 927],\n",
    "    [83, 91, 1, 645, 1253, 927],\n",
    "]\n",
    "\n",
    "padded_inputs = tf.keras.preprocessing.sequence.pad_sequences(raw_inputs,padding='post')\n",
    "padded_inputs"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "tf.Tensor(\n",
      "[[ True  True  True False False False]\n",
      " [ True  True  True  True  True False]\n",
      " [ True  True  True  True  True  True]], shape=(3, 6), dtype=bool)\n",
      "tf.Tensor(\n",
      "[[[7.110e+02 7.110e+02 7.110e+02 7.110e+02 7.110e+02 7.110e+02 7.110e+02\n",
      "   7.110e+02 7.110e+02 7.110e+02]\n",
      "  [6.320e+02 6.320e+02 6.320e+02 6.320e+02 6.320e+02 6.320e+02 6.320e+02\n",
      "   6.320e+02 6.320e+02 6.320e+02]\n",
      "  [7.100e+01 7.100e+01 7.100e+01 7.100e+01 7.100e+01 7.100e+01 7.100e+01\n",
      "   7.100e+01 7.100e+01 7.100e+01]\n",
      "  [0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00\n",
      "   0.000e+00 0.000e+00 0.000e+00]\n",
      "  [0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00\n",
      "   0.000e+00 0.000e+00 0.000e+00]\n",
      "  [0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00\n",
      "   0.000e+00 0.000e+00 0.000e+00]]\n",
      "\n",
      " [[7.300e+01 7.300e+01 7.300e+01 7.300e+01 7.300e+01 7.300e+01 7.300e+01\n",
      "   7.300e+01 7.300e+01 7.300e+01]\n",
      "  [8.000e+00 8.000e+00 8.000e+00 8.000e+00 8.000e+00 8.000e+00 8.000e+00\n",
      "   8.000e+00 8.000e+00 8.000e+00]\n",
      "  [3.215e+03 3.215e+03 3.215e+03 3.215e+03 3.215e+03 3.215e+03 3.215e+03\n",
      "   3.215e+03 3.215e+03 3.215e+03]\n",
      "  [5.500e+01 5.500e+01 5.500e+01 5.500e+01 5.500e+01 5.500e+01 5.500e+01\n",
      "   5.500e+01 5.500e+01 5.500e+01]\n",
      "  [9.270e+02 9.270e+02 9.270e+02 9.270e+02 9.270e+02 9.270e+02 9.270e+02\n",
      "   9.270e+02 9.270e+02 9.270e+02]\n",
      "  [0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00 0.000e+00\n",
      "   0.000e+00 0.000e+00 0.000e+00]]\n",
      "\n",
      " [[8.300e+01 8.300e+01 8.300e+01 8.300e+01 8.300e+01 8.300e+01 8.300e+01\n",
      "   8.300e+01 8.300e+01 8.300e+01]\n",
      "  [9.100e+01 9.100e+01 9.100e+01 9.100e+01 9.100e+01 9.100e+01 9.100e+01\n",
      "   9.100e+01 9.100e+01 9.100e+01]\n",
      "  [1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00\n",
      "   1.000e+00 1.000e+00 1.000e+00]\n",
      "  [6.450e+02 6.450e+02 6.450e+02 6.450e+02 6.450e+02 6.450e+02 6.450e+02\n",
      "   6.450e+02 6.450e+02 6.450e+02]\n",
      "  [1.253e+03 1.253e+03 1.253e+03 1.253e+03 1.253e+03 1.253e+03 1.253e+03\n",
      "   1.253e+03 1.253e+03 1.253e+03]\n",
      "  [9.270e+02 9.270e+02 9.270e+02 9.270e+02 9.270e+02 9.270e+02 9.270e+02\n",
      "   9.270e+02 9.270e+02 9.270e+02]]], shape=(3, 6, 10), dtype=float32)\n",
      "tf.Tensor(\n",
      "[[ True  True  True False False False]\n",
      " [ True  True  True  True  True False]\n",
      " [ True  True  True  True  True  True]], shape=(3, 6), dtype=bool)\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "from tensorflow.keras import layers\n",
    "embedding = layers.Embedding(input_dim=5000,output_dim=16,mask_zero=True)\n",
    "masked_output = embedding(padded_inputs)\n",
    "\n",
    "print(masked_output._keras_mask)\n",
    "\n",
    "masking_layer = layers.Masking()\n",
    "unmasked_embedding = tf.cast(tf.tile(tf.expand_dims(padded_inputs,axis=-1),[1,1,10]),tf.float32)\n",
    "print(unmasked_embedding)\n",
    "masked_embedding = masking_layer(unmasked_embedding)\n",
    "print(masked_embedding._keras_mask)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [],
   "source": [
    "from tensorflow import keras\n",
    "model = keras.Sequential([\n",
    "    layers.Embedding(input_dim=5000,output_dim=16,mask_zero=True),\n",
    "    layers.LSTM(32)\n",
    "])\n",
    "\n",
    "inputs = keras.Input(shape=(None,),dtype='int32')\n",
    "x = layers.Embedding(input_dim=5000,output_dim=16,mask_zero=True)(inputs)\n",
    "outputs = layers.LSTM(32)(x)\n",
    "\n",
    "model = keras.Model(inputs,outputs)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [
    {
     "data": {
      "text/plain": "<tf.Tensor: shape=(32, 32), dtype=float32, numpy=\narray([[ 0.0142168 ,  0.00534033,  0.01211633, ...,  0.00976348,\n         0.00793803,  0.01268861],\n       [ 0.0112685 ,  0.00734463,  0.00156741, ...,  0.00909687,\n         0.00525863, -0.00065796],\n       [ 0.00824948,  0.01320076,  0.01339381, ...,  0.00983908,\n         0.00304379,  0.00937503],\n       ...,\n       [-0.00789263, -0.00711612,  0.00054358, ..., -0.00158949,\n        -0.01145515, -0.01723501],\n       [ 0.00018029, -0.00867407, -0.00797173, ..., -0.00587268,\n         0.00060806, -0.01192909],\n       [ 0.00541659, -0.00638866,  0.00982996, ...,  0.00441375,\n        -0.00216445, -0.00633537]], dtype=float32)>"
     },
     "metadata": {},
     "output_type": "execute_result",
     "execution_count": 8
    }
   ],
   "source": [
    "import numpy as np\n",
    "class MyLayer(layers.Layer):\n",
    "    def __init__(self,**kwargs):\n",
    "        super(MyLayer,self).__init__(**kwargs)\n",
    "        self.embedding = layers.Embedding(input_dim=5000,output_dim=16,mask_zero=True)\n",
    "        self.lstm = layers.LSTM(32)\n",
    "    \n",
    "    def call(self,inputs):\n",
    "        x = self.embedding(inputs)\n",
    "        \n",
    "        mask = self.embedding.compute_mask(inputs)\n",
    "        output = self.lstm(x,mask=mask)\n",
    "        return output\n",
    "\n",
    "layer = MyLayer()\n",
    "x = np.random.random((32,10)) *100\n",
    "x = x.astype('int32')\n",
    "layer(x)\n",
    "\n",
    "    \n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "tf.Tensor(\n",
      "[[ True  True  True]\n",
      " [ True  True  True]\n",
      " [ True  True  True]], shape=(3, 3), dtype=bool)\n",
      "tf.Tensor(\n",
      "[[False False False]\n",
      " [ True  True False]\n",
      " [ True  True  True]], shape=(3, 3), dtype=bool)\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "class TemporalSplit(keras.layers.Layer):\n",
    "    \"\"\"Split the input tensor into 2 tensors along the time dimension.\"\"\"\n",
    "    def call(self,inputs):\n",
    "        return tf.split(inputs,2,axis=1)\n",
    "    \n",
    "    def compute_mask(self,inputs,mask=None):\n",
    "        if mask is None:\n",
    "            return None\n",
    "        return tf.split(mask,2,axis=1)\n",
    "\n",
    "first_half ,second_half = TemporalSplit()(masked_embedding)\n",
    "print(first_half._keras_mask)\n",
    "print(second_half._keras_mask)\n",
    "\n",
    "        \n",
    "    "
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "tf.Tensor(\n",
      "[[ True  True  True  True False  True  True  True  True  True]\n",
      " [ True  True  True  True  True False  True  True  True  True]\n",
      " [False  True  True  True  True  True  True  True  True  True]], shape=(3, 10), dtype=bool)\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "class CustomEmbedding(keras.layers.Layer):\n",
    "    def __init__(self,input_dim,output_dim,mask_zero=False,**kwargs):\n",
    "        super(CustomEmbedding,self).__init__(**kwargs)\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.mask_zero = mask_zero\n",
    "    \n",
    "    def build(self,input_shape):\n",
    "        self.embeddings = self.add_weight(\n",
    "            shape=(self.input_dim,self.output_dim),\n",
    "            initializer='random_normal',\n",
    "            dtype='float32'\n",
    "        )\n",
    "    \n",
    "    def call(self,inputs):\n",
    "        return tf.nn.embedding_lookup(self.embeddings,inputs)\n",
    "    \n",
    "    def compute_mask(self,inputs,mask=None):\n",
    "        if not self.mask_zero:\n",
    "            return None\n",
    "        return tf.not_equal(inputs,0)\n",
    "    \n",
    "layer = CustomEmbedding(10,32,mask_zero=True)\n",
    "x = np.random.random((3,10)) * 9\n",
    "x = x.astype('int32')\n",
    "\n",
    "y = layer(x)\n",
    "mask = layer.compute_mask(x)\n",
    "\n",
    "print(mask)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "outputs": [],
   "source": [
    "class MyActivation(keras.layers.Layer):\n",
    "    def __init__(self,**kwargs):\n",
    "        super(MyActivation,self).__init__(**kwargs)\n",
    "        self.supports_masking = True\n",
    "    \n",
    "    def call(self,inputs):\n",
    "        return tf.nn.relu(inputs)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "Mask found: Tensor(\"embedding_5/NotEqual:0\", shape=(None, None), dtype=bool)\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "inputs = keras.Input(shape=(None,),dtype='int32')\n",
    "x = layers.Embedding(input_dim = 5000,output_dim= 16,mask_zero=True)(inputs)\n",
    "x = MyActivation()(x)\n",
    "print(\"Mask found:\",x._keras_mask)\n",
    "outputs = layers.LSTM(32)(x)\n",
    "\n",
    "model = keras.Model(inputs,outputs)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "None\n",
      "Tensor(\"embedding_9/NotEqual:0\", shape=(None, None), dtype=bool)\n",
      "Tensor(\"embedding_9/NotEqual:0\", shape=(None, None), dtype=bool)\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "class TemporalSoftmax(keras.layers.Layer):\n",
    "    def call(self,inputs,mask=None):\n",
    "        broadcast_float_mask = tf.expand_dims(tf.cast(mask,'float32'),-1)\n",
    "        inputs_exp = tf.exp(inputs) * broadcast_float_mask\n",
    "        inputs_sum =tf.reduce_sum(inputs * broadcast_float_mask,axis=1,keepdims=True)\n",
    "        return inputs_exp /inputs_sum\n",
    "\n",
    "inputs = keras.Input(shape=(None,),dtype='int32')\n",
    "print(inputs._keras_mask)\n",
    "x = layers.Embedding(input_dim=10,output_dim=32,mask_zero=True)(inputs)\n",
    "print(x._keras_mask)\n",
    "x = layers.Dense(1)(x)\n",
    "print(x._keras_mask)\n",
    "outputs = TemporalSoftmax()(x)\n",
    "\n",
    "model = keras.Model(inputs,outputs)\n",
    "y = model(np.random.randint(0,10,size=(32,100)),np.random.random((32,100,1)))\n",
    "    "
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "conda-env-keras-py",
   "language": "python",
   "display_name": "Python [conda env:keras]"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}